# Convert RBR luppis v3 audio pkg to ztmz audio pkg

# python rbr_audiopkg_to_ztmz.py --rbr_csv='/mnt/c/Users/strawhat/ztmz/ztmz_pacenote/src/ZTMZ.PacenoteTool.RBR/Jannemod V3.csv' --db_rbr='/mnt/c/Users/strawhat/ztmz/ztmz_pacenote/src/ZTMZ.PacenoteTool.RBR/jannemod_v3.zdb' --db_ztmz='/mnt/c/Users/strawhat/ztmz/ztmz_pacenote/src/ZTMZ.PacenoteTool.Base/data.zdb' --rbr_audio_pkg='/mnt/c/Users/strawhat/Downloads/luppis V3 2024.10.24 老王中文语音1.069 4个版本/Luppis-v3-2024.8.31-老王中文语音1.069（原始版）/Plugins/Pacenote/sounds/Janne Laahanen EnNu B' --output='/mnt/c/Users/strawhat/Downloads/拉稀车手老王'

import asyncio
import pandas as pd
from argparse import ArgumentParser
import sqlite3
import glob
import os
from pathlib import Path
from tqdm import tqdm

async def parse_args():
    p = ArgumentParser()
    p.add_argument('--rbr_csv', help='path to luppis v3 json definition')
    p.add_argument('--rbr_audio_pkg', help='path to rbr audio pkg')
    p.add_argument('--db_ztmz', help='path to ztmz db')
    p.add_argument('--db_rbr', help='path to luppis v3 db')
    p.add_argument('--output', help='output path')
    args = p.parse_args()
    args = vars(args)
    return args

async def load_rbr_id_2_audio_files(rbr_csv):
    df = pd.read_csv(rbr_csv)
    id_2_audio_files = {}
    # iterrows
    for index, row in df.iterrows():
        id = row['Id']
        audio_files = row['Sounds']
        category = row['Category']
        if category == 'MODIFIER':
            continue
        if pd.isna(audio_files) or audio_files is None or audio_files == '':
            continue
        audio_files = audio_files.split(';')
        id_2_audio_files[id] = audio_files
    return id_2_audio_files

async def load_db_ztmz_id_2_filenames(db_ztmz):
    conn = sqlite3.connect(db_ztmz)
    c = conn.cursor()
    filenames = c.execute('SELECT * FROM pacenote_filenames')
    filenames = filenames.fetchall()
    conn.close()
    filenames_dict = {}
    for id, filename, is_primary in filenames:
        if filenames_dict.get(id) is None:
            filenames_dict[id] = []
        
        if is_primary:
            filenames_dict[id].insert(0, filename)
        else:
            filenames_dict[id].append(filename)

    return filenames_dict

async def load_db_rbr_id_2_ztmz_id(db_rbr):
    conn = sqlite3.connect(db_rbr)
    c = conn.cursor()
    pacenotes_2_ztmz = c.execute('SELECT * FROM pacenote_ztmz')
    pacenotes_2_ztmz = pacenotes_2_ztmz.fetchall()
    modifiers_2_ztmz = c.execute('SELECT * FROM modifier_ztmz')
    modifiers_2_ztmz = modifiers_2_ztmz.fetchall()
    conn.close()
    pacenote_id_2_ztmz_id = {}
    for id, ztmz_id, order_id in pacenotes_2_ztmz:
        if order_id == 0:
            pacenote_id_2_ztmz_id[id] = ztmz_id

    return pacenote_id_2_ztmz_id

async def load_rbr_audio_files(rbr_audio_pkg):
    # loop through the folder recursively to get all the ogg file paths
    audio_files = glob.glob(rbr_audio_pkg + '/**/*.ogg', recursive=True)
    audio_files_dict = { os.path.basename(f): f for f in audio_files }
    return audio_files_dict

async def fill_ztmz_audio_pkg(rbr_id_2_audio_files, ztmz_id_2_filenames, rbr_id_2_ztmz_id, rbr_audio_files, output):
    for rbr_id, ztmz_id in tqdm(rbr_id_2_ztmz_id.items()):
        # select destination folder
        filenames = ztmz_id_2_filenames[ztmz_id]
        selected_dest = None
        for filename in filenames:
            selected_dest = Path(output + '/' + filename)
            break

        if rbr_id not in rbr_id_2_audio_files:
            continue
        audio_files = rbr_id_2_audio_files[rbr_id]
        # gather audio_files from rbr_id_2_audio_files
        for audio_file in audio_files:
            if audio_file in rbr_audio_files:
                audio_file_path = rbr_audio_files[audio_file]
                audio_file_dest = selected_dest
                audio_file_dest = str(audio_file_dest)
                # copy audio file to the destination
                
                # mkdir
                if selected_dest is None:
                    print(f'No available destination for ztmz_id {ztmz_id}')
                    continue
                else:
                    selected_dest.mkdir(parents=True, exist_ok=True)
                os.system(f"cp '{audio_file_path}' '{audio_file_dest}'")
            else:
                print(f'Audio file {audio_file} not found in rbr_audio_files')

        
async def main():
    args = await parse_args()
    rbr_id_2_audio_files = await load_rbr_id_2_audio_files(args['rbr_csv'])
    ztmz_id_2_filenames = await load_db_ztmz_id_2_filenames(args['db_ztmz'])
    rbr_id_2_ztmz_id = await load_db_rbr_id_2_ztmz_id(args['db_rbr'])
    rbr_audio_files = await load_rbr_audio_files(args['rbr_audio_pkg'])
    await fill_ztmz_audio_pkg(rbr_id_2_audio_files, ztmz_id_2_filenames, rbr_id_2_ztmz_id, rbr_audio_files, args['output'])
    pass

if __name__ == "__main__":
    asyncio.run(main())
